package cn.ycc1.lianzi.controller;

import cn.ycc1.lianzi.service.ChatService;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.ApiException;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import io.reactivex.Flowable;
import jakarta.servlet.http.HttpServletRequest;
import lombok.Data;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * @author ycc
 * @date 2025/2/12
 */
@RestController
public class ChatController {
    @Autowired
    private ChatService chatService;

    // 存储对话上下文的伪代码（实际需用线程安全结构）
    private Map<String, List<ChatMessage>> sessions = new ConcurrentHashMap<>();

    @PostMapping(value = "/chat", consumes = MediaType.APPLICATION_JSON_VALUE,
            produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatStream(@RequestBody ChatRequest request,
                                 HttpServletRequest httpRequest) {
        SseEmitter emitter = new SseEmitter(180_000L); // 3分钟超时
        String sessionId = httpRequest.getSession().getId();

//        return emitter;
        AtomicInteger i = new AtomicInteger(1);

        try {
            Generation gen = new Generation();
            Message userMsg = Message.builder().role(Role.USER.getValue()).content("你是谁？").build();
            Flowable<GenerationResult> result = chatService.streamCallWithMessage(gen, userMsg);
            result.blockingForEach(message -> {
                String chunk = message.getOutput().getChoices().get(0).getMessage().getContent();
                    // 构建SSE消息
//                    emitter.send(SseEmitter.event()
//                            .data(chunk)
//                            .id(String.valueOf(i.incrementAndGet()))
//                            .name("chunk")
//                            .reconnectTime(5000));
//                System.out.println(chunk);
//                emitter.send(SseEmitter.event().data(chunk));
                try {
                    CompletableFuture.runAsync(() -> {
                        // 模拟从DashScope获取流式数据
                        while (!StringUtils.isEmpty(chunk)) {
                            try {
                                emitter.send(chunk);
                            } catch (IOException e) {
                                throw new RuntimeException(e);
                            }
                        }
                        emitter.complete();
                    });
                } catch (Exception ex) {
                    emitter.completeWithError(ex);
                }
            });
            emitter.complete();
        } catch (ApiException | NoApiKeyException | InputRequiredException e) {
            //logger.error("An exception occurred: {}", e.getMessage());
//            log.info("An exception occurred: {}", e.getMessage());
        }
        // System.exit(0);
        return emitter;
    }

    @PostMapping(value = "/chat2", consumes = MediaType.APPLICATION_JSON_VALUE,
            produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter chatStream2(@RequestBody ChatRequest request,
                                 HttpServletRequest httpRequest) {
        SseEmitter emitter = new SseEmitter(180_000L); // 3分钟超时
        String sessionId = httpRequest.getSession().getId();

        CompletableFuture.runAsync(() -> {
            try {
                // 获取或创建会话上下文
                List<ChatMessage> context = sessions.computeIfAbsent(sessionId,
                        k -> new CopyOnWriteArrayList<>());

                // 添加用户消息到上下文
                context.add(new ChatMessage("user", request.getMessage()));

                // 模拟大模型流式生成（实际调用AI API）
                for (int i = 1; i <= 10; i++) {
                    String chunk = "模型响应片段" + i + " ";

                    // 构建SSE消息
                    emitter.send(SseEmitter.event()
                            .data(chunk)
                            .id(String.valueOf(i))
                            .name("chunk")
                            .reconnectTime(5000));

                    Thread.sleep(500); // 模拟生成延迟
                }

                // 添加最终响应到上下文
                context.add(new ChatMessage("assistant", "完整响应内容"));
                emitter.complete();
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });

        return emitter;
    }

    // 数据结构定义
    static class ChatRequest {
        public String message;

        public String getMessage() {
            return message;
        }

        public void setMessage(String message) {
            this.message = message;
        }
    }

    @Data
    static class ChatMessage {
        String role;
        String content;

        // 构造方法...
        public ChatMessage(String assistant, String content) {
            this.role = assistant;
            this.content = content;
        }
    }
}
